/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package io.trino.jdbc.tuple.guide;

import com.google.common.collect.ImmutableMap;
import io.trino.jdbc.TrinoConnection;
import io.trino.jdbc.TrinoResultSet;
import io.trino.jdbc.TrinoStatement;
import io.trino.jdbc.tuple.guide.session.SessionPropType;
import io.trino.jdbc.tuple.rs.DataType;
import io.trino.jdbc.tuple.rs.column.ColumnDataParser;
import io.trino.jdbc.tuple.rs.column.ColumnFieldData;
import org.apache.commons.lang3.StringUtils;

import java.sql.*;
import java.util.*;
import java.util.concurrent.atomic.AtomicBoolean;

import static java.lang.String.format;

/**
 * @author fengyang(yueny09@163.com)
 * @date 2022年06月09日 下午3:33
 */
public final class Drivers {

    private Drivers(){
        //.
    }


    ///////////////////////////////////////////////////////////////////////
    //    连接驱动初始化             连接驱动初始化            连接驱动初始化    //
    ///////////////////////////////////////////////////////////////////////
    /**
     * 不负责关闭动作
     */
    public static TrinoConnection createConnection(String address)
            throws SQLException {
        return createConnection(address, "");
    }
    /**
     * 不负责关闭动作
     */
    public static TrinoConnection createConnection(String address, String extra)
            throws SQLException {
        return createConnection(address,  ImmutableMap.of(), extra);
    }

    /**
     * 不负责关闭动作
     */
    public static TrinoConnection createConnection(String address, Map<SessionPropType, String> session)
            throws SQLException {
        return createConnection(address, session, "");
    }

    /**
     * 不负责关闭动作
     */
    public static TrinoConnection createConnection(String address, Map<SessionPropType, String> session, String extra)
            throws SQLException {
        StringBuilder sb = new StringBuilder();
        sb.append("jdbc:trino://%s");

        // extra
        // sb.append("?useCursorFetch=true");
        if(StringUtils.isNotEmpty(extra)){
            // sb.append("&%s");
            sb.append("?%s");
        }

        String url = format(sb.toString(), address, StringUtils.isNotEmpty(extra) ? extra : "");

        // 不设置时区。 每个客户端展示的时区均为系统默认时区
//        // 设置时区
//        TimeZone.setDefault(TimeZone.getTimeZone(ZoneId.of(CommonConstants.DEFAULT_ZONE_ID)));

        try{
            Class.forName(CommonConstants.TRINO_DRIVER);
        }catch (ClassNotFoundException ex){
            ex.printStackTrace();
        }

        // trino 不会对密码进行验证， 可以为 null
        // 不会做用户名验证， 但不能为空
        TrinoConnection trinoConnection = (TrinoConnection) DriverManager.getConnection(url,
                CommonConstants.DEFAULT_USER_NAME, CommonConstants.DEFAULT_PASSWORD);

        if(session != null){
            session.forEach((k, v) -> {
                String key = CommonConstants.CATALOG_NAME_FOR_QUERY_ENGINE_DOT + k.getKey();

                trinoConnection.setSessionProperty(key, v);
            });
        }

        return trinoConnection;
    }

    /**
     * 不负责关闭动作
     */
    public static TrinoStatement createStatement(String address)
            throws SQLException {
        Connection connection = createConnection(address);

        return createStatement(connection);
    }

    /**
     * 不负责关闭动作
     */
    public static final TrinoStatement createStatement(String address, Map<SessionPropType, String> session)
            throws SQLException {
        Connection connection = createConnection(address, session);

        return createStatement(connection);
    }

    /**
     * 不负责关闭动作
     */
    public static final TrinoStatement createStatement(String address, Map<SessionPropType, String> session, String extra)
            throws SQLException {
        Connection connection = createConnection(address, session, extra);

        return createStatement(connection);
    }

    /**
     * 不负责关闭动作
     */
    public static TrinoStatement createStatement(Connection connection)
            throws SQLException {
        TrinoStatement statement = (TrinoStatement) connection.createStatement(ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY);
        createStatementBefore(statement);

        return statement;
    }

    /**
     * 不负责关闭动作
     *
     * @param resultSetType  eg. ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY
     */
    public static TrinoStatement createStatement(Connection connection, int resultSetType, int resultSetConcurrency)
            throws SQLException {
        TrinoStatement statement = (TrinoStatement) connection.createStatement(resultSetType, resultSetConcurrency);
        createStatementBefore(statement);

        return statement;
    }

    /**
     * 不负责关闭动作
     *
     * @param resultSetType  eg. ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY
     */
    public static TrinoStatement createStatement(Connection connection, int resultSetType, int resultSetConcurrency,
                                                 int resultSetHoldability)
            throws SQLException {
        TrinoStatement statement = (TrinoStatement) connection.createStatement(resultSetType,
                resultSetConcurrency, resultSetHoldability);
        createStatementBefore(statement);

        return statement;
    }

    private static final void createStatementBefore(TrinoStatement statement)
            throws SQLException {
        try{
//            statement.execute("USE ATOMDATA");

        }catch (Exception e){
            //.
        }
    }



    ///////////////////////////////////////////////////////////////////////
    //    转换结果集                 转换结果集                     转换结果集 //
    ///////////////////////////////////////////////////////////////////////

    /**
     * 转换所有结果集
     *
     * @param rs
     * @return  列表数据List， 包含行数据Map<String, ColumnData>和列信息 ColumnData
     *
     * @throws SQLException
     */
    public static List<Map<Integer, ColumnFieldData>> getRs(final TrinoResultSet rs) throws SQLException {
        // 获得结果集结构信息（元数据）
        ResultSetMetaData metadata = rs.getMetaData();

        int fetchSize;
        if(rs.getStatement().getFetchSize() > 0){
            fetchSize = rs.getStatement().getFetchSize();
        }else{
            fetchSize = Integer.MAX_VALUE;
        }

        AtomicBoolean hasPrinted = new AtomicBoolean(false);
        List<Map<Integer, ColumnFieldData>> list = new ArrayList<>();
        while (rs.next()) {
            Map<Integer, ColumnFieldData> rowData = getRowData(metadata, rs);
            if(!hasPrinted.get()){
                hasPrinted.set(true);
            }

            list.add(rowData);

            if(list.size() >= fetchSize){
                break;
            }
        }

        return list;
    }

    /**
     * 得到行数据
     *
     * @return  key is  position(>=1), value is ColumnData
     */
    public static Map<Integer, ColumnFieldData> getRowData(ResultSetMetaData metadata, ResultSet rs){
        Map<Integer, ColumnFieldData> rowMap = new TreeMap<>();

        try {
            for (int i = 1; i <= metadata.getColumnCount(); i++) {
                // rs.getObject(i) 获得当前行某一列字段的值
                String label = metadata.getColumnName(i);
                // columnType is java.sql.Types
                int columnType = metadata.getColumnType(i);
                ColumnFieldData columnData = getColumnData(label, columnType, rs.getObject(i), i);

                rowMap.put(columnData.getPosition(), columnData);
            }
        }catch (SQLException sqlException){
            sqlException.printStackTrace();
        }

        return Collections.unmodifiableMap(rowMap);
    }

    /**
     * 得到字段值， 包含具体类型
     *
     * @param columnType columnType is java.sql.Types
     */
    public static ColumnFieldData getColumnData(String label, int columnType, Object value, int position){
        ColumnFieldData.ColumnFieldDataBuilder builder = ColumnFieldData.builder();

        builder.label(StringUtils.lowerCase(label));
        builder.position(position);
        builder.columnType(columnType);

        Map.Entry<Object, DataType> entry = ColumnDataParser.parserDataType(value, columnType);
        if(entry != null){
            builder.dataType(entry.getValue());
            builder.value(entry.getKey());
        }else{
            // 暂不支持的类型
            builder.dataType(null);
            builder.value("N/A(NOT SUPPORT, columnType " + columnType + ")");
        }

        return builder.build();
    }

}
